import numpy as np
import scipy.io as scio
import matplotlib.pyplot as plt
import util
from util import moving_average
from util import moving_variance
from scipy.signal import find_peaks

def detect_peak(x, time_win, threshold_ratio=10):
    '''Detect peak via z-score. Not finished yet.'''

    x_ma = moving_average(x, time_win)
    x_std = np.sqrt(moving_variance(x, time_win))
    x_ma = np.roll(x_ma, time_win//2, axis=-1)
    x_std = np.roll(x_std, time_win//2, axis=-1)
    
    plt.plot(x_ma.squeeze())
    plt.savefig('/home/hepinjing/OCAI/experimental_process/temp/x_ma.png')
    plt.close()
    
    plt.plot(x_std.squeeze())
    plt.savefig('/home/hepinjing/OCAI/experimental_process/temp/x_std.png')
    plt.close()

    ret = np.zeros_like(x)
    res = x - x_ma

    scio.savemat('/home/hepinjing/OCAI/experimental_process/temp/peak_detection_middle_result.mat', {'x':x.reshape(-1,1), 'x_ma':x_ma.reshape(-1,1), 'x_std':x_std.reshape(-1,1), 'res':res.reshape(-1,1)})

    plt.plot(res.squeeze())
    plt.savefig('/home/hepinjing/OCAI/experimental_process/temp/res.png')
    plt.close()

    ret[res > threshold_ratio*x_std] = 1
    ret[res < -threshold_ratio*x_std] = -1

    return ret

if __name__ == '__main__':
    metric_data = scio.loadmat(
        '/home/hepinjing/OCAI/experimental_process/temp/metric.mat')
    metric = metric_data['metric']
    metric_abs = np.abs(metric)

    # metric_abs = np.array([*[0.0]*5, 1, *[0.0]*10, *[1.0]*10, *[0.0]*10])
    # metric_abs = metric_abs + 0.1 * np.random.randn(*metric_abs.shape)

    plt.plot(metric_abs.squeeze())
    plt.savefig('/home/hepinjing/OCAI/experimental_process/temp/metric_abs.png')
    plt.close()

    peak_detect_result = find_peaks(metric_abs.squeeze(), height=40, distance=128)
    # peak_detect_result = detect_peak(metric_abs.squeeze(), 101, 5)
    
    peak_candidate_indxes = np.where(peak_detect_result.squeeze()==1)[0]
    peak_indx = np.min(peak_candidate_indxes)

    print(peak_candidate_indxes[0:20])
    print(peak_indx)

    plt.plot(peak_detect_result.squeeze())
    plt.savefig('/home/hepinjing/OCAI/experimental_process/temp/peak_detect_sig.png')
    plt.close()

    scio.savemat('/home/hepinjing/OCAI/experimental_process/temp/peak_detection_result.mat', {'peak_detect_result':peak_detect_result})


    # a = np.stack([np.array([1.0,2.0,3.0,4.0,5.0,6.0]), np.array([0.0,1.0,2.0,3.0,4.0,5.0])], axis=0)
    # # a = a.reshape(1, -1)

    # b = util.moving_average(a, 3)

    # c = util.moving_variance(a, 3)

    # print(b)
    # print(c)



